import cv2
import numpy as np
from typing import List, Dict, Tuple
import logging

class ConnectionDetector:
    """连线识别模块"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
    
    def detect_connections(self, image: np.ndarray, binary_image: np.ndarray, nodes: List[Dict]) -> List[Dict]:
        """检测节点间的连接线"""
        self.logger.info("开始检测连接线")
        
        connections = []
        
        # 检测直线
        lines = self._detect_lines(binary_image)
        
        # 检测箭头
        arrows = self._detect_arrows(binary_image)
        
        # 建立连接关系
        connections = self._establish_connections(lines, arrows, nodes)
        
        self.logger.info(f"检测到 {len(connections)} 个连接")
        return connections
    
    def _detect_lines(self, binary_image: np.ndarray) -> List[Dict]:
        """检测图像中的直线"""
        lines = []
        
        # 使用霍夫变换检测直线
        hough_lines = cv2.HoughLinesP(binary_image, 1, np.pi/180, 
                                    threshold=50, minLineLength=30, maxLineGap=10)
        
        if hough_lines is not None:
            for line in hough_lines:
                x1, y1, x2, y2 = line[0]
                
                # 计算线段长度
                length = np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
                
                # 过滤太短的线段
                if length < 20:
                    continue
                
                line_info = {
                    'start': (x1, y1),
                    'end': (x2, y2),
                    'length': length,
                    'angle': np.arctan2(y2 - y1, x2 - x1)
                }
                lines.append(line_info)
        
        return lines
    
    def _detect_arrows(self, binary_image: np.ndarray) -> List[Dict]:
        """检测箭头"""
        arrows = []
        
        # 查找轮廓
        contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        for contour in contours:
            area = cv2.contourArea(contour)
            
            # 过滤太小的轮廓
            if area < 50 or area > 1000:
                continue
            
            # 计算轮廓的凸包
            hull = cv2.convexHull(contour)
            
            # 检查是否为箭头形状
            if self._is_arrow_shape(hull):
                # 计算箭头中心
                M = cv2.moments(contour)
                if M["m00"] != 0:
                    cx = int(M["m10"] / M["m00"])
                    cy = int(M["m01"] / M["m00"])
                    
                    arrow_info = {
                        'center': (cx, cy),
                        'contour': contour,
                        'direction': self._calculate_arrow_direction(hull)
                    }
                    arrows.append(arrow_info)
        
        return arrows
    
    def _is_arrow_shape(self, hull: np.ndarray) -> bool:
        """判断轮廓是否为箭头形状"""
        # 简化的箭头检测：检查凸包的点数和形状
        if len(hull) < 3 or len(hull) > 8:
            return False
        
        # 计算凸包的面积和周长
        area = cv2.contourArea(hull)
        perimeter = cv2.arcLength(hull, True)
        
        if perimeter == 0:
            return False
        
        # 计算紧凑度
        compactness = area / (perimeter * perimeter)
        
        # 箭头通常具有特定的紧凑度范围
        return 0.02 <= compactness <= 0.15
    
    def _calculate_arrow_direction(self, hull: np.ndarray) -> float:
        """计算箭头方向"""
        # 找到最长的边作为箭头的主方向
        max_length = 0
        direction = 0
        
        for i in range(len(hull)):
            p1 = hull[i][0]
            p2 = hull[(i + 1) % len(hull)][0]
            
            length = np.linalg.norm(p2 - p1)
            if length > max_length:
                max_length = length
                direction = np.arctan2(p2[1] - p1[1], p2[0] - p1[0])
        
        return direction
    
    def _establish_connections(self, lines: List[Dict], arrows: List[Dict], nodes: List[Dict]) -> List[Dict]:
        """建立节点间的连接关系"""
        connections = []
        
        # 为每个节点找到可能的连接
        for i, source_node in enumerate(nodes):
            source_pos = source_node['position']
            source_center = (source_pos['x'] + source_pos['width'] // 2,
                           source_pos['y'] + source_pos['height'] // 2)
            
            for j, target_node in enumerate(nodes):
                if i == j:
                    continue
                
                target_pos = target_node['position']
                target_center = (target_pos['x'] + target_pos['width'] // 2,
                               target_pos['y'] + target_pos['height'] // 2)
                
                # 检查是否有连接线连接这两个节点
                connection = self._find_connection_between_nodes(
                    source_center, target_center, lines, arrows
                )
                
                if connection:
                    connection['source'] = source_node['id']
                    connection['target'] = target_node['id']
                    connections.append(connection)
        
        return connections
    
    def _find_connection_between_nodes(self, source_center: Tuple[int, int], 
                                     target_center: Tuple[int, int],
                                     lines: List[Dict], arrows: List[Dict]) -> Dict:
        """查找两个节点间的连接"""
        # 计算节点间的距离
        distance = np.sqrt((target_center[0] - source_center[0])**2 + 
                          (target_center[1] - source_center[1])**2)
        
        # 查找连接这两个节点的线段
        connecting_lines = []
        for line in lines:
            # 检查线段是否连接两个节点
            if self._line_connects_points(line, source_center, target_center):
                connecting_lines.append(line)
        
        if not connecting_lines:
            return None
        
        # 选择最长的连接线
        best_line = max(connecting_lines, key=lambda x: x['length'])
        
        # 查找相关的箭头
        arrow = self._find_arrow_for_line(best_line, arrows)
        
        # 构建连接信息
        connection = {
            'type': 'direct',
            'points': [
                {'x': best_line['start'][0], 'y': best_line['start'][1]},
                {'x': best_line['end'][0], 'y': best_line['end'][1]}
            ],
            'length': best_line['length']
        }
        
        if arrow:
            connection['arrow'] = arrow['center']
            connection['direction'] = arrow['direction']
        
        return connection
    
    def _line_connects_points(self, line: Dict, point1: Tuple[int, int], 
                            point2: Tuple[int, int], tolerance: int = 20) -> bool:
        """检查线段是否连接两个点"""
        start = line['start']
        end = line['end']
        
        # 检查线段的端点是否接近给定的点
        start_to_p1 = np.sqrt((start[0] - point1[0])**2 + (start[1] - point1[1])**2)
        start_to_p2 = np.sqrt((start[0] - point2[0])**2 + (start[1] - point2[1])**2)
        end_to_p1 = np.sqrt((end[0] - point1[0])**2 + (end[1] - point1[1])**2)
        end_to_p2 = np.sqrt((end[0] - point2[0])**2 + (end[1] - point2[1])**2)
        
        # 一个端点接近point1，另一个端点接近point2
        return ((start_to_p1 < tolerance and end_to_p2 < tolerance) or
                (start_to_p2 < tolerance and end_to_p1 < tolerance))
    
    def _find_arrow_for_line(self, line: Dict, arrows: List[Dict], tolerance: int = 30) -> Dict:
        """为线段找到相关的箭头"""
        line_end = line['end']
        
        for arrow in arrows:
            arrow_center = arrow['center']
            distance = np.sqrt((arrow_center[0] - line_end[0])**2 + 
                             (arrow_center[1] - line_end[1])**2)
            
            if distance < tolerance:
                return arrow
        
        return None
    
    def detect_connections_with_nodes(self, image: np.ndarray, binary_image: np.ndarray, 
                                    nodes: List[Dict]) -> List[Dict]:
        """检测连接线并与节点关联"""
        connections = self.detect_connections(image, binary_image, nodes)
        
        # 过滤重复连接
        unique_connections = self._remove_duplicate_connections(connections)
        
        return unique_connections
    
    def _remove_duplicate_connections(self, connections: List[Dict]) -> List[Dict]:
        """移除重复的连接"""
        unique_connections = []
        seen_pairs = set()
        
        for connection in connections:
            pair = tuple(sorted([connection['source'], connection['target']]))
            
            if pair not in seen_pairs:
                seen_pairs.add(pair)
                unique_connections.append(connection)
        
        return unique_connections 